#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_modify_genome.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/22
@Description :
整理基因组序列文件，只保留常染色体（1-22）和性染色体（X,Y）。
将序列名修改为统一的 chr 格式（如 chr1、chrX），并按 chr1→chr22→chrX→chrY 排序输出。
"""

from argparse import ArgumentParser, Namespace
from Bio import SeqIO
from pathlib import Path
import re

# 定义排序顺序
CHROM_ORDER = {str(i): i for i in range(1, 23)}  # chr1-22
CHROM_ORDER.update({'X': 23, 'Y': 24})  # chrX=23, chrY=24


def chrom_sort_key(record):
    """
    根据 chr 名称返回排序键值，用于按 chr1-22,X,Y 排序
    """
    # 去掉前缀 'chr'
    chrom = record.id.replace("chr", "")
    return CHROM_ORDER.get(chrom, 100)  # 非常规染色体排到最后


def process_genome(infile: Path, outfile: Path):
    """
    读取输入 fasta 文件，筛选常染色体和性染色体序列，
    修改序列名为 chr 格式，按 chr1-22,X,Y 排序后写入输出文件。

    Args:
        infile (Path): 输入 fasta 文件路径
        outfile (Path): 输出 fasta 文件路径
    """
    # 确保输出目录存在
    outfile.parent.mkdir(parents=True, exist_ok=True)

    out_records = []

    with open(infile, "r", encoding="utf-8") as fr:
        records = SeqIO.parse(fr, "fasta")
        for record in records:
            # 匹配描述信息中的常染色体或性染色体（1-22, X, Y）
            chrom_match = re.search(r"chromosome\s+([1-9]|1[0-9]|2[0-2]|X|Y)$", record.description)
            if chrom_match:
                chrom = chrom_match.group(1)
                # 修改序列 id 和 description 为统一的 chr 格式
                record.id = f"chr{chrom}"
                record.description = f"chr{chrom}"
                out_records.append(record)

    # 按 chr1-22,X,Y 排序
    out_records.sort(key=chrom_sort_key)

    # 写入输出 fasta 文件
    with open(outfile, "w", encoding="utf-8") as fw:
        SeqIO.write(out_records, fw, "fasta")


def main(args: Namespace):
    infile = Path(args.infile).resolve()
    outfile = Path(args.outfile).resolve()
    process_genome(infile, outfile)


if __name__ == "__main__":
    parser = ArgumentParser(description="整理基因组序列文件，只保留常染色体（1-22）和性染色体（X,Y），并按顺序输出")
    parser.add_argument("-i", "--infile", type=str, required=True, help="输入 fasta 文件路径")
    parser.add_argument("-o", "--outfile", type=str, required=True, help="输出 fasta 文件路径")
    main(parser.parse_args())
